package constructdata.TestDataConstructors;

import codetree.CodeTree;
import codetree.CodeTreeOperation;
import codetree.TreeNode;
import treeview.DisplayTreeView;
import treeview.TreeView;

import java.io.*;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Construct Fine Grained Test Data
 * Created by wangxin on 2018/02/02.
 */
public class FineTestDataConstructor {

    private final boolean isDebug = false;

    // The marks
    private final String ENDMARK = "end";
    private final String CASEMARK = "case";
    private final String DEFAULTMARK = "default";
    private final String[] filterSigns = {"[","]"};
    private final String[] holeParentAdjust = {"else","elseif","catch","finally","case","default"};
    private final String[] specialLabel = {"end", "conditionEnd"};

    // The maximum hole number of the continuous holes
    private int MAX_HOLENUM = 5;
    // The construct trees
    private LinkedList<CodeTree> trees;
    // The predication holes correspondent to the constructTrees;
    private LinkedList<String> predictions;
    // The predicate node's complete class name
    private LinkedList<String> classnames;
    // The node who is the parent of the predicated one
    private LinkedList<TreeNode> parents;
    // CodeTreeOperator
    private CodeTreeOperation operator;
    // The size of hole
    private  LinkedList<String> holesizes;
    // The block predictions
    private LinkedList<String> blockpredictions;
    private String blockprediction;
    // The original statements
    private LinkedList<String> originalStatements;
    // The variable names
    private LinkedList<String> variableNames;
    // The block original statements
    private LinkedList<String> blockstatements;
    private String blockstatement;

    // test
    private int count = 0;

    // testcase trace
    private LinkedList<String> testcaseTraces;
    // testcase pool: <trace, testcase>
    private Map<String, String> testcasePool;
    private Map<Integer,String> sourceLines;
    private String trace;
    private String testcaseDir;

    /**
     * Constructor
     * */
    public FineTestDataConstructor(){
        operator = new CodeTreeOperation();
        this.count  = 0;
    }

    public void construct(CodeTree tree,
                          FileWriter treeWriter, FileWriter predictionWriter, FileWriter classWriter, FileWriter generationNodeWriter, FileWriter treeSentenceWriter, FileWriter jarWriter, FileWriter holeSizeWriter,
                          FileWriter traceWriter, // trace back
                          FileWriter blockpredictionWriter, // block of predictions (more lines)
                          FileWriter originalStatementsWriter, // original statements
                          FileWriter variableNamesWriter,// variable names
                          FileWriter linesWriter,// lines writer
                          FileWriter testCaseTraceWriter,
                          FileWriter blockOriginalStatementWriter,// block of original statements
                          String testcaseDir,
                          boolean isCompleteFlag){
        this.testcaseDir = testcaseDir;

        LinkedList<LinkedList> result = getConstructTrainingData(tree);
        List<CodeTree> treeList = result.get(0);
        List<String> predictionList = result.get(1);
        List<String> classList = result.get(2);
        List<TreeNode> generationNodeList = result.get(3);
        List<String> holeSizeList = result.get(4);// record size of hole
        List<String> blockpredictionList = result.get(5);

        List<String> originalStatementsList = result.get(6);
        List<String> variableNameList = result.get(7);
        List<String> testCaseTraceList = result.get(8); // -- testcase trace
        List<String> blockOriginalStatementList = result.get(9);


        //FindJarHandler findJarHandler = new FindJarHandler();
        for (int i = 0; i < treeList.size(); i++) {
            try {
                //String jar = findJarHandler.getPackage(classList.get(i));
                //if(jar != null) {
                CodeTree tempTree = treeList.get(i);
                operator.saveRegularizedTreeInFile(operator.regularization(tempTree), treeWriter);
                operator.saveTrainingPredictionInFile(predictionList.get(i), predictionWriter);
                operator.saveTrainingPredictionInFile(classList.get(i), classWriter);
                //operator.saveTrainingPredictionInFile(jar, jarWriter);
                int parentnum = generationNodeList.get(i).getSerialNumber();
                operator.saveTrainingPredictionInFile(parentnum + " " + ((parentnum != 0) ? generationNodeList.get(i).getCompleteMethodDeclaration() : ""), generationNodeWriter);
                operator.saveTreeStringFormatInFile(tempTree, treeSentenceWriter, isCompleteFlag);

                operator.saveTrainingPredictionInFile(holeSizeList.get(i), holeSizeWriter);

                operator.saveTrainingPredictionInFile(tree.getFunctionTrace(), traceWriter);
                operator.saveTrainingPredictionInFile(blockpredictionList.get(i), blockpredictionWriter);

                operator.saveTrainingPredictionInFile(originalStatementsList.get(i), originalStatementsWriter);
                operator.saveTrainingPredictionInFile(variableNameList.get(i), variableNamesWriter);

                operator.saveTrainingPredictionInFile((tempTree.getLines(tempTree.getRoot()) + 1) +"", linesWriter);

                operator.saveTrainingPredictionInFile( testCaseTraceList.get(i),testCaseTraceWriter);

                operator.saveTrainingPredictionInFile(blockOriginalStatementList.get(i), blockOriginalStatementWriter);
                //}
            } catch (Exception e) {
                System.err.println(e.getMessage());
            } catch (Error e){
                System.err.println(e.getMessage());
            }
        }
    }

    /**
     * Construct training tree from the input code tree
     * Return the LinkedList with 6 list:
     * 1. The training trees with some continuous holes
     * 2. The predication holes correspondent to the training trees
     * 3. The class names of the prediction
     * 4. The parent node of the prediction
     * 5. The holesize
     * 6. The block preditions with multi-lines
     *
     * @param completeTree: the complete code tree
     * */
    public LinkedList<LinkedList> getConstructTrainingData(CodeTree completeTree){
        // Init
        trees = new LinkedList<>();
        predictions = new LinkedList<>();
        classnames = new LinkedList<>();
        parents = new LinkedList<>();
        holesizes = new LinkedList<>();
        blockpredictions = new LinkedList<>();
        originalStatements = new LinkedList<>();
        variableNames = new LinkedList<>();
        testcaseTraces = new LinkedList<>();
        testcasePool = new HashMap<>(); // -- for test case
        sourceLines = new HashMap<>();
        blockstatements = new LinkedList<>();

        // Construct
        //MAX_HOLENUM = completeTree.getTotalNumber(); // set max hole number
        for (int i = 1; i <= MAX_HOLENUM; i++) {
            construct(completeTree, 1, i);// the serial number of root is 1.
        }

        // Return
        LinkedList<LinkedList> result = new LinkedList<>();
        result.addLast(trees);
        result.addLast(predictions);
        result.addLast(classnames);
        result.addLast(parents);
        result.addLast(holesizes);
        result.addLast(blockpredictions);
        result.addLast(originalStatements);
        result.addLast(variableNames);

        result.addLast(testcaseTraces);

        result.addLast(blockstatements);
        return result;
    }

    /**
     * @param completeTree: training data from this tree
     * @param serialNumber: the first node to be remove is with this serial number
     * @param holeNumber: the number of continuous holes is holeNumber
     * */
    private void construct(CodeTree completeTree, int serialNumber, int holeNumber){

        // Make the replica of the completeTree to avoid destroying.
        CodeTree tree;
        tree = copyCodeTree(completeTree);
        operator.setSerialNumberofEachNode(tree);

        // Initialize the blockprediction & blockstatement
        blockprediction  = "";
        blockstatement = "";

        // The source code
        trace = Util.getTrace(completeTree);
        sourceLines = Util.getSourceLines(trace);

        if(holeNumber > 0 && tree.getTotalNumber() >= serialNumber){
            boolean isSuccess = true;
            List predicts = null;
            TreeNode node =  tree.getTreeNode(serialNumber);
            TreeNode constrainParent = null;
            if(node != null){
                constrainParent = node.getParentNode();
            }
            int loopSerialNumber = serialNumber;
            int i = 0;
            for (i = 0; i < holeNumber; i++) {
                if(tree.getTotalNumber() == 1){// only root node.
                    isSuccess = false;
                    break;
                }
                List temp = null;
                if(i == 0){
                    if((predicts = remove(tree,constrainParent,loopSerialNumber)) == null
                            ){
                        isSuccess = false;
                        break;
                    }
                }
                else if((temp = remove(tree,constrainParent,loopSerialNumber)) == null){
                    isSuccess = false;
                    break;
                }
                // Consider the condition case will change the serial number!(remove end, -1)
                if(i == 0){
                    loopSerialNumber = Integer.parseInt(predicts.get(3).toString());
                }
                else if (temp != null){
                    loopSerialNumber = Integer.parseInt(temp.get(3).toString());
                }
            }
            if(isSuccess){
                String testcase = Util.convert2SourceCode(sourceLines);
                String predict = (String) predicts.get(0);
                if((!testcasePool.containsValue(testcase)) || isSpecialLabel(predict)) { // real, success when testcase pool do not contains the test case
                    // For debug
                    count++;
                    if (isDebug) {
                        System.out.println(count + ". " + predicts + " parentnum: " + ((TreeNode) predicts.get(2)).getSerialNumber() + " ..hole: " + holeNumber);
                    }
                    predictions.addLast(predict);
                    trees.addLast(tree);
                    classnames.addLast((String) predicts.get(1));
                    TreeNode generationNode = (TreeNode) predicts.get(2);
                    parents.addLast(generationNode);

                    // Add hole node
                    addHole(tree, generationNode);
                    holesizes.addLast("" + holeNumber);

                    // block predictions
                    if (isDebug) {
                        System.out.println("block prediction: " + blockprediction.trim());// -- debug
                        //displayTree(tree, true, blockprediction.trim());
                    }
                    blockpredictions.addLast(blockprediction.trim());
                    blockstatements.addLast(blockstatement.trim());
                    blockstatement = "";

                    // statements
                    originalStatements.addLast((String) predicts.get(4));
                    // variable names
                    variableNames.addLast((String) predicts.get(5));

                    // extract the function
                    String methodinfo = tree.getMethodInfo();
                    int b = Integer.parseInt(methodinfo.split(" ")[0]);
                    int e = Integer.parseInt(methodinfo.split(" ")[1]);
                    String functionblock = Util.convert2SourceCode(sourceLines, b, e);

                    // construct the testcase: testcase = sourceinfo + functionblock + }
                    String sourceinfo = tree.getSourceInfo();
                    testcase = sourceinfo + functionblock + "}";

                    // test case save
                    String testcaseTrace = saveTestCase(testcase, count, trace, b + "to" + e, "FineGrain");
                    testcasePool.put(testcaseTrace,testcase);
                    // test case trace
                    testcaseTraces.addLast(testcaseTrace);
                }
            }
            serialNumber++;
            construct(completeTree, serialNumber, holeNumber);
        }
    }

    private String saveTestCase(String testcase, int count, String trace, String lines, String flag) {
        String path = testcaseDir + trace.replace("/", "_")+"_"+lines+"_"+count+"_"+flag+".java";
        try {
            BufferedWriter bufferedWriter = new BufferedWriter(new FileWriter(path));
            bufferedWriter.write(testcase);
            bufferedWriter.flush();
            bufferedWriter.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return path;
    }

    private CodeTree copyCodeTree(CodeTree completeTree) {
        try {
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            ObjectOutputStream oos = new ObjectOutputStream(bos);
            oos.writeObject(completeTree);
            ByteArrayInputStream bis = new ByteArrayInputStream(bos.toByteArray());
            ObjectInputStream ois = new ObjectInputStream(bis);
            CodeTree tree = (CodeTree) ois.readObject();
            return tree;
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        }
        return null;
    }

    private void addHole(CodeTree tree, TreeNode generationNode) {
        TreeNode hole = new TreeNode();
        hole.setClassName("hole");
        hole.setCompleteClassName("hole");
        hole.setMethodName("");
        hole.setCompleteMethodName("");
        hole.setAddMethodName(false);
        tree.addNode(tree.getTreeNode(generationNode.getSerialNumber()), hole);
        operator.setSerialNumberofEachNode(tree);
    }

    /**
     * Remove the node with serialNumber in the code tree
     * Set its parent as its children's parent
     * Return its complete method declaration, classname, and the parent node
     *
     * @param tree: the code tree
     * @param constrainParent: the parent of this node
     * @param serialNumber: the serial number of this node
     * */
    private LinkedList<Object> remove(CodeTree tree, TreeNode constrainParent, int serialNumber){
        LinkedList<Object> result = new LinkedList<>();

        TreeNode node =  tree.getTreeNode(serialNumber);
        if(node == null || node.isCondition()){// Not exist
            return null;
        }

        TreeNode parent = node.getParentNode();

        if(parent != null && !parent.equals(constrainParent)){
            return null;// Stop remove when the parent is not the constrain one
        }

        List<TreeNode> children = node.getChildNodes();
        String label = node.getCompleteMethodDeclaration();
        String className = node.getCompleteClassName();

        String statement = node.getStatement();
        List<String> previousVariableNames = node.getPreviousVariableNames();
        String variablename = previousVariableNames.size()>0?previousVariableNames.get(0):"";
        for (int i = 1; i < previousVariableNames.size(); i++) {
            variablename += " " + previousVariableNames.get(i);
        }

        // When removing a node, add the block predictions
        blockprediction += getBlockPredictions(node);
        constructTestCase(node,false);

        // leaf node / inner node with end mark
        if (children.isEmpty() || children.get(0).getCompleteMethodDeclaration().equals(ENDMARK)) {
            if (parent != null) {
                parent.getChildNodes().remove(node);
            }
        }
        else { // root with children / inner node
            // Control node
            if (node.isControl()) {// get the successor to append to the parent
                children = new ArrayList<TreeNode>();
                TreeNode successor = getSuccessor(node);
                if(successor != null){
                    children.add(successor);
                }
            }

            // set parent
            for (TreeNode child : children) {
                child.setParentNode(parent);
            }

            // set children
            if (parent != null) {
                // the index of node
                int index = parent.getChildNodes().indexOf(node);
                if (children.size() > 0) { // transplant
                    // transplant the child to the parent
                    parent.getChildNodes().set(index, children.get(0));
                }
                else{// moving-brothers
                    for (int i = index; i < parent.getChildNodes().size() - 1; i++) {
                        parent.getChildNodes().set(i, parent.getChildNodes().get(i + 1));
                    }
                    parent.getChildNodes().remove(parent.getChildNodes().get(parent.getChildNodes().size() - 1));
                }
            } else {// root
                if(children.size() > 0 && !children.get(0).getCompleteMethodDeclaration().equals("else") && !children.get(0).getCompleteMethodDeclaration().equals("elseif")) {
                    tree.setRoot(children.get(0));
                }
                else{
                    tree.setRoot(new TreeNode());
                    return null;// the tree is not exist now.
                }
            }
        }

        // Reorder nodes
        operator.setSerialNumberofEachNode(tree);
        result.addLast(label);
        result.addLast(filterSigns(className));

        // Reassign hole parent
        if(isReassignHoleParent(label)){
            parent = node.getHoleParentNode();
        }

        if(parent == null) {
            parent = new TreeNode();
            parent.setSerialNumber(0);// the root do not have a parent
        }
        result.addLast(parent);
        result.addLast(serialNumber);// set the serial number

        result.addLast(statement);
        result.addLast(variablename);
        // result:
        // 0: label
        // 1: class
        // 2: parent
        // 3: serialNumber
        // 4: statement
        // 5: variable
        return result;
    }

    private boolean isSpecialLabel(String label){
        List<String> l = Arrays.stream(specialLabel).collect(Collectors.toList());
        return l.contains(label);
    }

    // Replace the begin~end lines as holes to construct testcase.
    private void constructTestCase(TreeNode node, boolean isThrough) {
        String label = node.getCompleteMethodDeclaration();
        String scinfo = node.getInfo();
        if(scinfo == null || node.isCondition()){
            return;
        }

        int begin = Util.getBeginLine(node);
        int end = Util.getEndLine(node);
        String stmt = Util.getStmt(node);
//        System.out.println(begin);
//        System.out.println(end);
        if((!node.isControl()) && begin==end) {
            blockstatement = blockstatement + " " + Util.replaceStmt(begin, stmt, sourceLines, " /*hole*/ ");
        }
        else {
            blockstatement = blockstatement + " " + Util.replaceLine2Line(begin, end, sourceLines, " /*hole*/ ");
        }

        if(node.isControl()){
            if(label.equals("if") && node.getChildNodes().size() > 1){
                if(node.getChildNodes().size() > 2) {// else/ elseif
                    TreeNode secondChildNode = node.getChildNodes().get(2);
                    String secondChildLabel = secondChildNode.getCompleteMethodDeclaration();
                    if (secondChildLabel.equals("else") || secondChildLabel.equals("elseif")) {
                        constructTestCase(secondChildNode,true);
                    }
                }
            }
            else if(isThrough && label.equals("elseif")){
                if(node.getChildNodes().size() == 3){// "successor" part - else/elseif
                    constructTestCase(node.getChildNodes().get(2),true);
                }
            }
            else if (label.equals("else") || label.equals("elseif") || label.equals("for") || label.equals("while") || label.equals("foreach") || label.equals("doWhile")) {
                // replace all...
            }
            else if (label.equals("switch")) {
                // replace all...
            }
            else if (label.equals("try")) {
                List<TreeNode> children = node.getChildNodes();
                if (children.size() > 0) {
                    for (int i = 0; i < children.size(); i++) {
                        String tmpLabel = children.get(i).getCompleteMethodDeclaration();
                        if(tmpLabel.equals("catch") || tmpLabel.equals("finally")){
                            constructTestCase(children.get(i),false);
                        }
                    }
                }
            }
            else if (label.equals("catch")){
                // replace all...
            }
            else if(label.equals("finally")){
                // replace all...
            }
        }
    }

    // Get multi-lines of predictions when removing this node
    private String getBlockPredictions(TreeNode node) {
        String buff = "";

        // record label of the first node
        String label = node.getCompleteMethodDeclaration();
        buff += " " + label;

        // take the nodes inner control-node into consideration
        if(node.isControl()){
            // append body/then part
            if(label.equals("if") && node.getChildNodes().size() > 1){
                buff += " " + node2String(node.getChildNodes().get(0)); // condition
                buff += " " + node2String(node.getChildNodes().get(1)); // then
                if(node.getChildNodes().size() > 2) {
                    String secondChildLabel = node.getChildNodes().get(2).getCompleteMethodDeclaration();
                    if (secondChildLabel.equals("else") || secondChildLabel.equals("elseif")) {
                        buff += " " + node2String(node.getChildNodes().get(2));
                    }
                }
            }
            else if(label.equals("else")){
                if(node.getChildNodes().size() > 0) {
                    buff += " " + node2String(node.getChildNodes().get(0));
                }
            }
            else if(label.equals("elseif")){
                buff += " " + node2String(node.getChildNodes().get(0)); // condition
                if(node.getChildNodes().size() > 1) {
                    buff += " " + node2String(node.getChildNodes().get(1));
                }
            }
            else if (label.equals("for") || label.equals("while") || label.equals("foreach") || label.equals("doWhile")) {
                buff += " " + node2String(node.getChildNodes().get(0)); // condition
                if(node.getChildNodes().size() > 1) {// body
                    buff += " " + node2String(node.getChildNodes().get(1));
                }
            }
            else if (label.equals("switch")){
                buff += " " + node2String(node.getChildNodes().get(0)); // condition
                for (int i = 0; i < node.getChildNodes().size(); i++) {
                    String childLabel = node.getChildNodes().get(i).getCompleteMethodDeclaration();
                    if(childLabel.equals(CASEMARK) || childLabel.equals(DEFAULTMARK)){
                        buff += " " + node2String(node.getChildNodes().get(i));
                    }
                }
            }
            else if (label.equals("case") || label.equals("default")){
                if(node.getChildNodes().size() > 0) {
                    buff += " " + node2String(node.getChildNodes().get(0));
                }
            }
            else if (label.equals("try")){ // try (xxx) catch xxx catch xxx finally xxx (successor)
                for (int i = 0; i < node.getChildNodes().size(); i++) {
                    String childLabel = node.getChildNodes().get(i).getCompleteMethodDeclaration();
                    if(childLabel.equals("catch") || childLabel.equals("finally")){
                        buff += " " + node2String(node.getChildNodes().get(i));
                    }
                    else if(i != (node.getChildNodes().size() - 1)){
                        buff += " " + node2String(node.getChildNodes().get(i));
                    }
                }
            }
            else if (label.equals("catch")){// catch xxx
                if(node.getChildNodes().size() > 0) {
                    buff += " " + node2String(node.getChildNodes().get(0));
                }
            }
            else if (label.equals("finally")){// finally xxx
                if(node.getChildNodes().size() > 0) {
                    buff += " " + node2String(node.getChildNodes().get(0));
                }
            }
        }
        return buff;
    }

    // DFS node to string
    private String node2String(TreeNode treeNode) {
        String buff = treeNode.getCompleteMethodDeclaration();
        for (int i = 0; i < treeNode.getChildNodes().size(); i++) {
            buff += " " + node2String(treeNode.getChildNodes().get(i));
        }
        return buff;
    }

    private TreeNode getSuccessor(TreeNode node){
        String label = node.getCompleteMethodDeclaration();
        List<TreeNode> children = node.getChildNodes();
        if(label.equals("if")){
            if(children.size() == 4) {// There exist successor part
                return children.get(3);
            }
            else if(children.size() == 3){
                String childLabel = children.get(2).getCompleteClassName();
                if(childLabel.equals("else") || childLabel.equals("elseif")){// "false" part
                    return null;// no successor part
                }
                else{// There exist successor part
                    return children.get(2);
                }
            }
        }
//        else if(label.equals("else")){
//            return null; // "else" no successor
//        }
        else if(label.equals("elseif")){
            if(children.size() == 3){// "successor" part - else/elseif
                return children.get(2);
            }
        }
        else if (label.equals("for") || label.equals("while") || label.equals("foreach") || label.equals("doWhile") ) {
            if (children.size() == 3) {// There exist successor part
                return children.get(2);
            }
        }
        else if (label.equals("switch")) {
            if (children.size() > 1) {
                String tmpLabel = children.get(children.size() - 1).getCompleteMethodDeclaration();
                if (!tmpLabel.equals(DEFAULTMARK) && !tmpLabel.equals(CASEMARK)) {
                    return children.get(children.size() - 1);
                }
            }
        }
        else if (label.equals("try")) {
            if (children.size() > 0) {
                String tmpLabel = children.get(children.size() - 1).getCompleteMethodDeclaration();
                if(!tmpLabel.equals("catch") && !tmpLabel.equals("finally")){
                    return children.get(children.size() - 1);
                }
            }
        }
        else if (label.equals("catch")){
            if(children.size() == 2) {
                return children.get(1);// "successor" part - catch/finally
            }
        }
//        else if(label.equals("finally")){
//            return null;
//        }
        return null;
    }

    // filter the signs
    private String filterSigns(String string){
        for (String sign: filterSigns) {
            string = string.replace(sign, "");
        }
        return string;
    }

    // Reassign hole parent.
    private boolean isReassignHoleParent(String label){
        for(String l : holeParentAdjust){
            if(l.equals(label)){
                return true;
            }
        }
        return false;
    }

    // for testing ...
    private void displayTree(CodeTree codeTree, boolean isCompleteFlag, String title) {
        TreeView treeView = new TreeView();
        treeView.convertCodeTree(codeTree, isCompleteFlag);
        new DisplayTreeView(treeView.getTree(),title);
    }
}
